import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import rasterio
from rasterio.plot import show
from scipy.interpolate import interp1d

# 设置输入文件夹路径和输出文件夹路径
input_folder = "./data/month"  # 替换为你的输入文件夹
output_folder = "./output"          # 替换为你的输出文件夹

# 确保输出文件夹存在
os.makedirs(output_folder, exist_ok=True)

# 月份列表（4月到10月）
months = [4, 5, 6, 7, 8, 9, 10]

# 读取所有月度NDVI数据
ndvi_data = []
for month in months:
    file_path = os.path.join(input_folder, f"NDVI_M{month}.tif")  # 假设文件名格式为 ndvi_4.tif, ndvi_5.tif 等
    with rasterio.open(file_path) as src:
        ndvi_data.append(src.read(1))  # 假设是单波段数据

# 转换为NumPy数组
ndvi_array = np.array(ndvi_data)

# 检查是否有缺失值（NaN）
has_nan = np.isnan(ndvi_array).any()
if has_nan:
    print("存在缺失值，正在处理...")
    # 对每个像元进行插值填充缺失值
    for i in range(ndvi_array.shape[1]):
        for j in range(ndvi_array.shape[2]):
            # 获取该像元的时间序列
            pixel_values = ndvi_array[:, i, j]
            # 找到非NaN的值
            valid_indices = ~np.isnan(pixel_values)
            if valid_indices.any():
                # 如果有有效值，用线性插值填充缺失值
                f = interp1d(
                    np.where(valid_indices)[0],
                    pixel_values[valid_indices],
                    kind='linear',
                    fill_value="extrapolate"
                )
                ndvi_array[:, i, j] = f(np.arange(len(pixel_values)))
            else:
                # 如果所有值都是NaN，用0填充（或根据需求调整）
                ndvi_array[:, i, j] = 0

# 原始时间点（月度）
original_time = np.array(months)

# 插值后的时间点（每半个月）
interpolated_time = np.linspace(4, 10, 14)  # 4到10月，每半个月一个点

# 对每个像元进行插值
interpolated_ndvi = np.zeros((len(interpolated_time), ndvi_array.shape[1], ndvi_array.shape[2]))

for i in range(ndvi_array.shape[1]):
    for j in range(ndvi_array.shape[2]):
        # 获取该像元的时间序列
        pixel_values = ndvi_array[:, i, j]
        # 进行插值
        f = interp1d(original_time, pixel_values, kind='linear', fill_value="extrapolate")
        interpolated_ndvi[:, i, j] = f(interpolated_time)

# 保存插值后的结果
for idx, time_point in enumerate(interpolated_time):
    output_file = os.path.join(output_folder, f"ndvi_{time_point:.2f}.tif")
    with rasterio.open(file_path) as src:  # 使用第一个文件的元数据
        meta = src.meta.copy()
        meta.update({"count": 1})
    with rasterio.open(output_file, "w", **meta) as dst:
        dst.write(interpolated_ndvi[idx, :, :], 1)

# 绘制插值前后的曲线
# 选择一个像元（例如中心点）进行可视化
center_row = ndvi_array.shape[1] // 2
center_col = ndvi_array.shape[2] // 2

# 原始数据
original_values = ndvi_array[:, center_row, center_col]
# 插值后的数据
interpolated_values = interpolated_ndvi[:, center_row, center_col]

# 绘图
plt.figure(figsize=(12, 6))

# 原始数据曲线
plt.subplot(1, 2, 1)
plt.plot(original_time, original_values, 'o-', label='Original Monthly NDVI')
plt.title('Original Monthly NDVI')
plt.xlabel('Month')
plt.ylabel('NDVI')
plt.legend()
plt.grid(True)

# 插值后的数据曲线
plt.subplot(1, 2, 2)
plt.plot(interpolated_time, interpolated_values, 'o-', label='Interpolated Biweekly NDVI')
plt.title('Interpolated Biweekly NDVI')
plt.xlabel('Month')
plt.ylabel('NDVI')
plt.legend()
plt.grid(True)

plt.tight_layout()
plt.show()